

# from stereo.utils.Filter.Filter_disp import fill_disp_near_keypoints
from stereo.utils.stereoRectify_process import stereo_rectify
from stereo.utils.Filter.OneEuroFilter_python import OneEuroFilter, EMAFilter
from stereo.POSE_process import PoseEstimationWorker, BODY_PARTS
from multiprocessing import Process, Event, Queue
from stereo.SGBM_process_crop import StereoMatcherWorker
import numpy as np
import time
import cv2
import traceback
import atexit
EMAFilter_ins = EMAFilter()
from stereo.utils.config import Config
 


class StereoProcessor:
    def __init__(self, pose_task_queue, pose_result_queue, pose_event,
                       depth_task_queue, depth_result_queue, depth_event):
        self.pose_worker = PoseEstimationWorker(pose_task_queue, pose_result_queue, pose_event)
        self.pose_worker.start()
        
        self.depth_worker = StereoMatcherWorker(depth_task_queue, depth_result_queue, depth_event)
        self.depth_worker.start()

    # 一个进程不断存入depth和pose队列图像，然后进行处理，一个进程不断存入depth和pose队列图像，然后进行处理
    def process_frame(self, left_image, right_image):
        self.depth_task_queue.put((left_image, right_image))
        self.pose_task_queue.put(left_image)

    # 一个进程不断的从depth和pose取同帧数据
    def get_frame(self):
        while True:
            if not self.depth_result_queue.empty() and not self.pose_result_queue.empty():
                depth_image, pose_image = self.depth_result_queue.get(), self.pose_result_queue.get()
                return depth_image, pose_image
            else:
                time.sleep(0.01)
                continue

# 一个进程不断存入depth和pose队列图像，然后进行处理，一个进程不断存入depth和pose队列图像，然后进行处理
def put_frame(pose_queue, depth_queue):
    # try:
    job_id = 0
    t1 = time.time()
    frames = 0
    FPS = 30
    interval = 1.0 / FPS
    left_image = np.random.randint(0, 255, (480, 640), dtype=np.uint8)
    right_image = np.random.randint(0, 255, (480, 640), dtype=np.uint8)
    while True:
        pose_queue.put((left_image, job_id))
        depth_queue.put((left_image, right_image, job_id))
        frames += 1
        job_id += 1

        # 每隔1秒打印一次FPS
        if time.time() - t1 >= 1.0:
            print(f"[输入] FPS: {frames} 帧/秒, 队列大小: {pose_queue.qsize()}, {depth_queue.qsize()}")
            frames = 0
            t1 = time.time()
        time.sleep(interval)  # 模拟 30FPS 摄像头
    # except Exception as e:
    #     traceback.print_exc()

def get_frame(pose_result_queue, depth_result_queue):
    start_time = time.time()
    last_frame_time = time.time()
    frame_count = 0
    total_time = 0.0

    while True:
        if not pose_result_queue.empty() and not depth_result_queue.empty():
            kps_dict_2d, kps_2d_list, job_id_pose = pose_result_queue.get()
            disp, job_id_depth = depth_result_queue.get()

            now = time.time()
            elapsed = now - last_frame_time
            total_time += elapsed
            frame_count += 1
            last_frame_time = now

            if now - start_time >= 1.0:
                fps = frame_count / (now - start_time)
                avg_time = total_time / frame_count if frame_count > 0 else 0
                print(f"[输出] FPS: {fps:.2f}, 平均耗时: {avg_time*1000:.2f} ms, 队列大小: {pose_result_queue.qsize()}, {depth_result_queue.qsize()}, 帧ID: {job_id_pose}, {job_id_depth}")
                start_time = now
                total_time = 0
                frame_count = 0
        else:
            time.sleep(0.005)

import psutil
def terminate_all():
    current = psutil.Process()
    children = current.children(recursive=True)
    for child in children:
        print(f"正在终止子进程: {child.pid}")
        child.terminate()
        child.wait()


if __name__ == "__main__":
    sender_flag = False
    try:

        # 创建共享资源
        pose_task_queue = Queue()
        pose_result_queue = Queue()
        pose_event = Event()

        depth_task_queue = Queue()
        depth_result_queue = Queue()
        depth_event = Event()

        # 启动推理子进程
        processor = StereoProcessor(
            pose_task_queue, pose_result_queue, pose_event,
            depth_task_queue, depth_result_queue, depth_event
        )

        # 生产者/消费者进程
        put_process = Process(target=put_frame, args=(pose_task_queue, depth_task_queue))
        get_process = Process(target=get_frame, args=(pose_result_queue, depth_result_queue))

        put_process.start()
        get_process.start()

        put_process.join()
        get_process.join()

    except KeyboardInterrupt:
        print("检测到 Ctrl+C，准备退出...")
        atexit.register(terminate_all)  # ✅ 一定放在主逻辑前
